import gradio as gr
import numpy as np
import pickle
from PIL import Image, ImageOps

# 加载模型函数
def load_knn_model():
    try:
        with open("best_knn_model.pkl", "rb") as f:
            return pickle.load(f)
    except FileNotFoundError:
        print("错误：模型文件不存在，请先运行optimal_knn.py生成")
        exit(1)

# 图像处理函数集合
def extract_image_data(sketch):
    if not sketch:
        raise ValueError("请绘制数字后再识别")
    if isinstance(sketch, dict):
        valid_keys = ["composite", "background", "foreground"]
        for k in valid_keys:
            if k in sketch:
                return np.array(sketch[k])
    return np.array(sketch)

def to_grayscale(image_array):
    if image_array.ndim == 3:
        return np.uint8(0.2 * image_array[:, :, 0] + 
                       0.6 * image_array[:, :, 1] + 
                       0.2 * image_array[:, :, 2])
    return np.uint8(image_array)

def preprocess_image(sketch):
    img = extract_image_data(sketch)
    if img.ndim not in (2, 3):
        raise ValueError(f"图像维度无效: {img.ndim}")
    gray = to_grayscale(img)
    pil_img = Image.fromarray(gray)
    resized = pil_img.resize((8, 8), Image.Resampling.LANCZOS)
    inverted = ImageOps.invert(resized)
    normalized = np.array(inverted) / 16
    return normalized.flatten().reshape(1, -1)

# 预测函数
def recognize_digit(sketch, model):
    try:
        processed = preprocess_image(sketch)
        prediction = model.predict(processed)[0]
        probabilities = model.predict_proba(processed)[0]
        prob_result = {str(i): f"{p:.4f}" for i, p in enumerate(probabilities)}
        return f"识别结果: {prediction}", prob_result
    except Exception as e:
        return f"识别失败: {str(e)}", {}

# 构建界面
def build_interface(model):
    with gr.Blocks(title="KNN手写数字识别") as demo:
        gr.Markdown("## 手写数字识别应用")
        gr.Markdown("在画板中绘制0-9的数字，点击识别按钮查看结果")
        
        with gr.Row():
            with gr.Column(scale=1):
                sketch_pad = gr.Sketchpad(height=300, width=300, label="数字绘制区")
                with gr.Row():
                    predict_btn = gr.Button("识别数字", variant="primary")
                    clear_btn = gr.Button("清空")
            with gr.Column(scale=1):
                result_text = gr.Textbox(label="识别结果")
                prob_label = gr.Label(label="概率分布")
        
        predict_btn.click(
            fn=lambda s: recognize_digit(s, model),
            inputs=sketch_pad,
            outputs=[result_text, prob_label]
        )
        
        clear_btn.click(
            fn=lambda: (None, "", {}),
            inputs=[],
            outputs=[sketch_pad, result_text, prob_label]
        )
    
    return demo

# 主入口
if __name__ == "__main__":
    model = load_knn_model()
    print("模型加载完成")
    app = build_interface(model)
    print("应用启动，访问 http://localhost:2200 使用")
    app.launch(
        server_name="0.0.0.0",
        server_port=2200,
        show_error=True
    )